#!/bin/bash

# AIQuant 系统状态检查脚本
# 用于检查AIQuant量化投研系统的运行状态

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_header() {
    echo -e "${CYAN}$1${NC}"
}

# 检查服务状态
check_service_status() {
    local service_name=$1
    local port=$2
    local endpoint=$3
    
    if lsof -Pi :$port -sTCP:LISTEN -t >/dev/null 2>&1; then
        if curl -f -s "$endpoint" > /dev/null 2>&1; then
            log_success "$service_name 运行正常 (端口: $port)"
            return 0
        else
            log_warning "$service_name 端口开放但服务异常 (端口: $port)"
            return 1
        fi
    else
        log_error "$service_name 未运行 (端口: $port)"
        return 1
    fi
}

# 检查进程状态
check_process_status() {
    local pid_file=$1
    local service_name=$2
    
    if [ -f "$pid_file" ]; then
        local pid=$(cat "$pid_file")
        if ps -p $pid > /dev/null 2>&1; then
            local cpu_usage=$(ps -p $pid -o %cpu --no-headers | tr -d ' ')
            local mem_usage=$(ps -p $pid -o %mem --no-headers | tr -d ' ')
            local start_time=$(ps -p $pid -o lstart --no-headers)
            log_success "$service_name 运行中 (PID: $pid, CPU: ${cpu_usage}%, MEM: ${mem_usage}%)"
            echo "    启动时间: $start_time"
            return 0
        else
            log_error "$service_name PID文件存在但进程未运行 (PID: $pid)"
            return 1
        fi
    else
        log_warning "$service_name PID文件不存在"
        return 1
    fi
}

# 检查数据库连接
check_database() {
    log_info "检查数据库连接..."
    
    python3 -c "
import sys
sys.path.append('.')
try:
    from aiquant.database.connection import DatabaseManager
    db = DatabaseManager()
    db.test_connection()
    print('✓ 数据库连接正常')
except ImportError:
    print('⚠ 无法导入数据库模块')
    sys.exit(1)
except Exception as e:
    print(f'✗ 数据库连接失败: {e}')
    sys.exit(1)
" 2>/dev/null
    
    if [ $? -eq 0 ]; then
        log_success "数据库连接正常"
        return 0
    else
        log_error "数据库连接失败"
        return 1
    fi
}

# 检查Redis连接
check_redis() {
    log_info "检查Redis连接..."
    
    if command -v redis-cli &> /dev/null; then
        if redis-cli ping > /dev/null 2>&1; then
            local redis_info=$(redis-cli info server | grep "redis_version" | cut -d: -f2 | tr -d '\r')
            log_success "Redis连接正常 (版本: $redis_info)"
            return 0
        else
            log_error "Redis连接失败"
            return 1
        fi
    else
        log_warning "Redis CLI未安装，无法检查连接"
        return 1
    fi
}

# 检查系统资源
check_system_resources() {
    log_header "=== 系统资源状态 ==="
    
    # CPU使用率
    local cpu_usage=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | cut -d'%' -f1)
    echo "CPU使用率: ${cpu_usage}%"
    
    # 内存使用率
    local mem_info=$(free | grep Mem)
    local total_mem=$(echo $mem_info | awk '{print $2}')
    local used_mem=$(echo $mem_info | awk '{print $3}')
    local mem_usage=$(echo "scale=1; $used_mem * 100 / $total_mem" | bc)
    echo "内存使用率: ${mem_usage}%"
    
    # 磁盘使用率
    echo "磁盘使用率:"
    df -h | grep -E '^/dev/' | awk '{print "  " $1 ": " $5 " (" $3 "/" $2 ")"}'
    
    # 负载平均值
    local load_avg=$(uptime | awk -F'load average:' '{print $2}')
    echo "系统负载:$load_avg"
}

# 检查网络连接
check_network() {
    log_header "=== 网络连接状态 ==="
    
    # 检查监听端口
    echo "监听端口:"
    netstat -tlnp 2>/dev/null | grep -E ':(8000|8501|5432|6379|3000|9090)' | while read line; do
        echo "  $line"
    done
    
    # 检查外部连接
    echo ""
    echo "外部连接测试:"
    
    # 测试互联网连接
    if ping -c 1 8.8.8.8 > /dev/null 2>&1; then
        echo "  ✓ 互联网连接正常"
    else
        echo "  ✗ 互联网连接异常"
    fi
    
    # 测试DNS解析
    if nslookup google.com > /dev/null 2>&1; then
        echo "  ✓ DNS解析正常"
    else
        echo "  ✗ DNS解析异常"
    fi
}

# 检查日志文件
check_logs() {
    log_header "=== 日志文件状态 ==="
    
    if [ -d "logs" ]; then
        echo "日志文件:"
        for log_file in logs/*.log; do
            if [ -f "$log_file" ]; then
                local size=$(du -h "$log_file" | cut -f1)
                local lines=$(wc -l < "$log_file")
                echo "  $(basename $log_file): $size ($lines 行)"
                
                # 检查最近的错误
                local errors=$(tail -100 "$log_file" | grep -i error | wc -l)
                if [ $errors -gt 0 ]; then
                    echo "    ⚠ 最近100行中有 $errors 个错误"
                fi
            fi
        done
    else
        echo "日志目录不存在"
    fi
}

# 检查配置文件
check_config() {
    log_header "=== 配置文件状态 ==="
    
    # 检查主要配置文件
    local config_files=("requirements.txt" "config.yaml" ".env")
    
    for config_file in "${config_files[@]}"; do
        if [ -f "$config_file" ]; then
            echo "✓ $config_file 存在"
        else
            echo "✗ $config_file 不存在"
        fi
    done
    
    # 检查环境变量
    echo ""
    echo "重要环境变量:"
    echo "  ENVIRONMENT: ${ENVIRONMENT:-未设置}"
    echo "  PYTHONPATH: ${PYTHONPATH:-未设置}"
    echo "  DATABASE_URL: ${DATABASE_URL:-未设置}"
    echo "  REDIS_URL: ${REDIS_URL:-未设置}"
}

# 生成状态报告
generate_report() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    local report_file="logs/status_report_$(date '+%Y%m%d_%H%M%S').txt"
    
    if [ ! -d "logs" ]; then
        mkdir -p logs
    fi
    
    {
        echo "AIQuant系统状态报告"
        echo "生成时间: $timestamp"
        echo "=================================="
        echo ""
        
        # 重新运行所有检查并输出到文件
        main --quiet
        
    } > "$report_file"
    
    log_success "状态报告已生成: $report_file"
}

# 主函数
main() {
    local quiet_mode=false
    
    if [ "$1" = "--quiet" ]; then
        quiet_mode=true
    fi
    
    if [ "$quiet_mode" = false ]; then
        log_header "=================================="
        log_header "    AIQuant 系统状态检查"
        log_header "=================================="
        echo ""
    fi
    
    # 服务状态检查
    log_header "=== 核心服务状态 ==="
    
    local api_status=0
    local web_status=0
    local db_status=0
    local redis_status=0
    
    # 检查API服务
    check_service_status "API服务" "8000" "http://localhost:8000/health" || api_status=1
    check_process_status ".api.pid" "API进程" || true
    
    echo ""
    
    # 检查Web服务
    check_service_status "Web服务" "8501" "http://localhost:8501" || web_status=1
    check_process_status ".web.pid" "Web进程" || true
    
    echo ""
    
    # 检查数据库
    check_database || db_status=1
    
    echo ""
    
    # 检查Redis
    check_redis || redis_status=1
    
    echo ""
    
    # 后台任务状态
    log_header "=== 后台任务状态 ==="
    check_process_status ".data_collector.pid" "数据收集器" || true
    check_process_status ".market_monitor.pid" "市场监控器" || true
    check_process_status ".anomaly_detector.pid" "异常检测器" || true
    
    echo ""
    
    if [ "$quiet_mode" = false ]; then
        # 系统资源检查
        check_system_resources
        echo ""
        
        # 网络连接检查
        check_network
        echo ""
        
        # 日志文件检查
        check_logs
        echo ""
        
        # 配置文件检查
        check_config
        echo ""
    fi
    
    # 总体状态
    log_header "=== 总体状态 ==="
    
    local total_issues=$((api_status + web_status + db_status + redis_status))
    
    if [ $total_issues -eq 0 ]; then
        log_success "所有核心服务运行正常"
    else
        log_warning "发现 $total_issues 个问题需要关注"
    fi
    
    echo ""
    echo "访问地址:"
    echo "  Web界面: http://localhost:8501"
    echo "  API文档: http://localhost:8000/docs"
    echo "  健康检查: http://localhost:8000/health"
    
    if [ "$quiet_mode" = false ]; then
        echo ""
        echo "使用 './scripts/start.sh' 启动系统"
        echo "使用 './scripts/stop.sh' 停止系统"
        echo "使用 '$0 --report' 生成详细报告"
    fi
}

# 处理参数
case "${1:-}" in
    --help|-h)
        echo "AIQuant系统状态检查脚本"
        echo "用法: $0 [选项]"
        echo "选项:"
        echo "  --help, -h     显示帮助信息"
        echo "  --report       生成状态报告"
        echo "  --quiet        静默模式"
        exit 0
        ;;
    --report)
        generate_report
        exit 0
        ;;
    --quiet)
        main --quiet
        ;;
    "")
        main
        ;;
    *)
        log_error "未知参数: $1"
        echo "使用 --help 查看帮助信息"
        exit 1
        ;;
esac